package com.zbkj.client.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.Map;
import java.util.List;
import java.util.ArrayList;
import java.util.Iterator;

/**
 * WebSocket会话管理器
 * 用于管理用户连接和推送消息
 */
@Slf4j
@Component
public class WebSocketSessionManager {
    
    // 存储用户ID到WebSocket会话列表的映射，支持同一用户多个连接
    private final Map<Integer, List<WebSocketSessionInfo>> userSessions = new ConcurrentHashMap<>();
    
    // 定时清理过期连接的调度器
    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
    
    // WebSocket连接有效期（30分钟）
    private static final long SESSION_TIMEOUT = 30 * 60 * 1000; // 30分钟
    
    /**
     * WebSocket会话信息类
     */
    private static class WebSocketSessionInfo {
        private final WebSocketSession session;
        private volatile long lastActiveTime;
        
        public WebSocketSessionInfo(WebSocketSession session) {
            this.session = session;
            this.lastActiveTime = System.currentTimeMillis();
        }
        
        public WebSocketSession getSession() {
            return session;
        }
        
        public long getLastActiveTime() {
            return lastActiveTime;
        }
        
        public void updateLastActiveTime() {
            this.lastActiveTime = System.currentTimeMillis();
        }
        
        public boolean isExpired() {
            return System.currentTimeMillis() - lastActiveTime > SESSION_TIMEOUT;
        }
    }
    
    public WebSocketSessionManager() {
        // 启动定时清理任务，每5分钟执行一次
        scheduler.scheduleAtFixedRate(this::cleanExpiredSessions, 5, 5, TimeUnit.MINUTES);
    }
    
    /**
     * 注册用户会话
     * @param userId 用户ID
     * @param session WebSocket会话
     */
    public void registerSession(Integer userId, WebSocketSession session) {
        userSessions.computeIfAbsent(userId, k -> new ArrayList<>())
                .add(new WebSocketSessionInfo(session));
        log.info("用户 {} 的WebSocket会话已注册，会话ID: {}", userId, session.getId());
    }
    
    /**
     * 移除用户的指定会话
     * @param userId 用户ID
     * @param session WebSocket会话
     */
    public void removeSession(Integer userId, WebSocketSession session) {
        List<WebSocketSessionInfo> sessions = userSessions.get(userId);
        if (sessions != null) {
            sessions.removeIf(info -> info.getSession().getId().equals(session.getId()));
            if (sessions.isEmpty()) {
                userSessions.remove(userId);
            }
            log.info("用户 {} 的WebSocket会话已移除，会话ID: {}", userId, session.getId());
        }
    }
    
    /**
     * 移除用户的所有会话
     * @param userId 用户ID
     */
    public void removeAllSessions(Integer userId) {
        userSessions.remove(userId);
        log.info("用户 {} 的所有WebSocket会话已移除", userId);
    }
    
    /**
     * 刷新指定会话的活跃时间
     * @param session WebSocket会话
     */
    public void refreshSocket(WebSocketSession session) {
        Integer userId = (Integer) session.getAttributes().get("userId");
        if (userId != null) {
            List<WebSocketSessionInfo> sessions = userSessions.get(userId);
            if (sessions != null) {
                for (WebSocketSessionInfo info : sessions) {
                    if (info.getSession().getId().equals(session.getId())) {
                        info.updateLastActiveTime();
                        log.debug("刷新用户 {} 的WebSocket会话活跃时间，会话ID: {}", userId, session.getId());
                        break;
                    }
                }
            }
        }
    }
    
    /**
     * 向指定用户推送消息（轮询所有连接）
     * @param userId 用户ID
     * @param message 消息内容
     * @return 成功推送的连接数
     */
    public int sendMessageToUser(Integer userId, String message) {
        List<WebSocketSessionInfo> sessions = userSessions.get(userId);
        if (sessions == null || sessions.isEmpty()) {
            log.warn("用户 {} 没有活跃的WebSocket连接", userId);
            return 0;
        }
        
        int successCount = 0;
        Iterator<WebSocketSessionInfo> iterator = sessions.iterator();
        
        while (iterator.hasNext()) {
            WebSocketSessionInfo info = iterator.next();
            WebSocketSession session = info.getSession();
            
            if (!session.isOpen()) {
                iterator.remove();
                log.info("移除用户 {} 的已关闭连接，会话ID: {}", userId, session.getId());
                continue;
            }
            
            try {
                session.sendMessage(new TextMessage(message));
                successCount++;
                log.debug("向用户 {} 推送消息成功，会话ID: {}", userId, session.getId());
            } catch (IOException e) {
                log.error("向用户 {} 推送消息失败，会话ID: {}，错误: {}", userId, session.getId(), e.getMessage());
                iterator.remove();
            }
        }
        
        // 如果用户没有任何活跃连接，清理该用户记录
        if (sessions.isEmpty()) {
            userSessions.remove(userId);
        }
        
        log.info("向用户 {} 推送消息完成，成功推送到 {} 个连接", userId, successCount);
        return successCount;
    }
    
    /**
     * 检查用户是否有活跃的WebSocket连接
     * @param userId 用户ID
     * @return 是否有活跃连接
     */
    public boolean hasActiveSession(Integer userId) {
        List<WebSocketSessionInfo> sessions = userSessions.get(userId);
        if (sessions == null || sessions.isEmpty()) {
            return false;
        }
        
        // 清理已关闭的连接
        sessions.removeIf(info -> !info.getSession().isOpen());
        
        if (sessions.isEmpty()) {
            userSessions.remove(userId);
            return false;
        }
        
        return true;
    }
    
    /**
     * 获取当前活跃连接数
     * @return 活跃连接数
     */
    public int getActiveSessionCount() {
        return userSessions.values().stream()
                .mapToInt(sessions -> (int) sessions.stream()
                        .filter(info -> info.getSession().isOpen())
                        .count())
                .sum();
    }
    
    /**
     * 清理过期的会话连接
     */
    private void cleanExpiredSessions() {
        log.debug("开始清理过期的WebSocket连接");
        int cleanedCount = 0;
        
        Iterator<Map.Entry<Integer, List<WebSocketSessionInfo>>> userIterator = userSessions.entrySet().iterator();
        
        while (userIterator.hasNext()) {
            Map.Entry<Integer, List<WebSocketSessionInfo>> entry = userIterator.next();
            Integer userId = entry.getKey();
            List<WebSocketSessionInfo> sessions = entry.getValue();
            
            Iterator<WebSocketSessionInfo> sessionIterator = sessions.iterator();
            while (sessionIterator.hasNext()) {
                WebSocketSessionInfo info = sessionIterator.next();
                
                if (info.isExpired() || !info.getSession().isOpen()) {
                    try {
                        if (info.getSession().isOpen()) {
                            info.getSession().close();
                        }
                    } catch (IOException e) {
                        log.warn("关闭过期WebSocket连接时发生异常: {}", e.getMessage());
                    }
                    
                    sessionIterator.remove();
                    cleanedCount++;
                    log.info("清理用户 {} 的过期WebSocket连接，会话ID: {}", userId, info.getSession().getId());
                }
            }
            
            // 如果用户没有任何连接，移除用户记录
            if (sessions.isEmpty()) {
                userIterator.remove();
            }
        }
        
        if (cleanedCount > 0) {
            log.info("清理过期WebSocket连接完成，共清理 {} 个连接", cleanedCount);
        }
    }
} 